import asyncio
import json
import time
import os
import re
import signal
import sys
import serial
from websockets import connect, exceptions
from smbus2 import SMBus
from configparser import ConfigParser
from threading import Thread

# -------------------------- 基础配置 --------------------------
CONFIG_FILE = "/www/wwwroot/new-human-control/pi/config.ini"
PCA9685_ADDR_0 = 0x40
PCA9685_ADDR_1 = 0x41
PWM_FREQ = 50
board_name = "body"
wifi_configured = False
running = True
config = ConfigParser()

# -------------------------- 串口配置 --------------------------
SERIAL_CONFIG = {
    "port": "/dev/ttyAMA0",
    "baudrate": 115200,
    "bytesize": serial.EIGHTBITS,
    "parity": serial.PARITY_NONE,
    "stopbits": serial.STOPBITS_ONE,
    "timeout": 0.01,
}
ser = None

# -------------------------- PCA9685类 --------------------------
class PCA9685:
    def __init__(self, address, bus_number=1):
        self.address = address
        self.bus = SMBus(bus_number)
        try:
            data = self.bus.read_byte_data(self.address, 0x00)
            print(f"PCA9685 检测成功 (0x{self.address:02X}): 0x{data:02X}")
        except OSError as e:
            raise OSError(f"PCA9685 连接失败 (0x{self.address:02X}): {e}")
        self.software_reset()
        self.set_pwm_freq(PWM_FREQ)

    def software_reset(self):
        self.bus.write_byte_data(self.address, 0x06, 0x01)
        time.sleep(0.1)

    def set_pwm_freq(self, freq):
        prescaleval = 25000000.0 / 4096.0 / float(freq) - 1.0
        prescale = int(round(prescaleval))
        old_mode = self.bus.read_byte_data(self.address, 0x00)
        new_mode = (old_mode & 0x7F) | 0x10
        self.bus.write_byte_data(self.address, 0x00, new_mode)
        self.bus.write_byte_data(self.address, 0xFE, prescale)
        self.bus.write_byte_data(self.address, 0x00, old_mode)
        time.sleep(0.005)
        self.bus.write_byte_data(self.address, 0x00, old_mode | 0x80)

    def set_pwm(self, channel, on, off):
        if not (0 <= channel <= 15):
            raise ValueError("通道必须0-15")
        if not (0 <= on <= 4095 and 0 <= off <= 4095):
            raise ValueError("ON/OFF必须0-4095")
        self.bus.write_byte_data(self.address, 0x06 + 4*channel, on&0xFF)
        self.bus.write_byte_data(self.address, 0x07 + 4*channel, on>>8)
        self.bus.write_byte_data(self.address, 0x08 + 4*channel, off&0xFF)
        self.bus.write_byte_data(self.address, 0x09 + 4*channel, off>>8)

    def set_servo_angle(self, channel, angle):
        pulse = int(205 + (angle/180.0)*205)
        self.set_pwm(channel, 0, pulse)

# -------------------------- 初始化PCA9685 --------------------------
try:
    pwm_0x40 = PCA9685(PCA9685_ADDR_0)
    pwm_0x41 = PCA9685(PCA9685_ADDR_1)
except OSError as e:
    print(f"PCA9685初始化失败: {e}")
    sys.exit(1)

# -------------------------- 通道映射函数 --------------------------
def get_pwm_device_and_channel(board_id, channel): 
    try:
        board_id = int(board_id)
    except ValueError:
        print(f"无效板子ID: {board_id}")
        return None, None
    if board_id not in (0, 1):
        print(f"板子ID超范围: {board_id}")
        return None, None

    try:
        physical_channel = int(channel)
    except ValueError:
        print(f"无效通道: {channel}")
        return None, None
    if not (0 <= physical_channel <= 15):
        print(f"通道超范围: {physical_channel}")
        return None, None

    if board_id == 0:
        return pwm_0x40, physical_channel
    elif board_id == 1:
        return pwm_0x41, physical_channel

# -------------------------- 舵机控制 --------------------------
def handle_servo_control(obj):
    print(f"[舵机控制] 收到指令：{obj}")
    if "b" not in obj or "c" not in obj or "p" not in obj:
        print("缺少必要字段：'b'、'c'或'p'")
        return
    try:
        board_id = obj["b"]
        channel = obj["c"]
        pulse = int(obj["p"])

        pwm_device, physical_channel = get_pwm_device_and_channel(board_id, channel)
        if not pwm_device or physical_channel is None:
            print(f"[舵机控制] 设备/通道映射失败：board_id={board_id}, channel={channel}")
            return 

        pulse = max(100, min(599, pulse))
        print(f"[舵机控制] 执行：板子ID {board_id}（0x{pwm_device.address:02X}） 通道{channel} 脉宽{pulse}")
        pwm_device.set_pwm(physical_channel, 0, pulse)

    except ValueError as e:
        print(f"[舵机控制] 参数错误: {e}")
    except Exception as e:
        print(f"[舵机控制] 执行失败: {e}")

# -------------------------- 文件监听逻辑 --------------------------
CMD_FILE_PATH = "/tmp/servo_cmd.json"
last_file_mtime = 0

def read_cmd_file():
    global last_file_mtime
    try:
        file_stat = os.stat(CMD_FILE_PATH)
        current_mtime = file_stat.st_mtime
        if current_mtime <= last_file_mtime:
            return
        last_file_mtime = current_mtime

        with open(CMD_FILE_PATH, "r", encoding="utf-8") as f:
            content = f.read().strip()
        if not content:
            print("指令文件为空")
            return

        payload = json.loads(content)
        print(f"从文件读取指令: {payload}")
        if isinstance(payload, list):
            for servo in payload:
                if isinstance(servo, dict):
                    handle_servo_control(servo)
        elif isinstance(payload, dict):
            handle_servo_control(payload)
        else:
            print("指令格式错误")

    except FileNotFoundError:
        pass
    except json.JSONDecodeError as e:
        print(f"指令文件JSON解析错误: {e}")
    except Exception as e:
        print(f"读取指令文件失败: {e}")

def file_listener_thread():
    print("文件监听线程已启动（检测路径：{}）".format(CMD_FILE_PATH))
    while running:
        read_cmd_file()
        time.sleep(0.5)

# -------------------------- 串口初始化 --------------------------
def init_serial():
    global ser
    try:
        ser = serial.Serial(**SERIAL_CONFIG)
        if ser.is_open:
            print(f"串口初始化成功：{ser.port}（波特率：{ser.baudrate}）")
            return True
        else:
            print("串口初始化失败：端口未打开")
            return False
    except serial.SerialException as e:
        print(f"串口错误：{str(e)}")
        return False

# -------------------------- 串口数据读取+解析（增强鲁棒性） --------------------------
def read_and_parse_serial():
    global ser
    if not ser or not ser.is_open:
        return None
    while ser.in_waiting:
        data = ser.readline().decode("ascii", errors="ignore").strip()
        if not data:
            continue
        print(f"[Server 串口原始数据] {data}")
        # 强制匹配格式：必须是 "数字_数字_数字"
        if re.match(r"^\d+_\d+_\d+$", data):
            parts = data.split("_")
            try:
                board_id = int(parts[0])
                channel = int(parts[1])
                angle = int(parts[2])
                cmd = {"b": board_id, "c": channel, "p": angle}
                print(f"[Server 串口解析结果] 生成指令: {cmd}")
                return cmd
            except ValueError:
                print(f"[Server 串口错误] 数字转换失败：{data}")
        else:
            print(f"[Server 串口错误] 格式不匹配（需数字_数字_数字）：{data}")
    return None

# -------------------------- 串口监听线程（新增） --------------------------
def serial_listener_thread():
    print("串口监听线程已启动")
    while running:
        serial_cmd = read_and_parse_serial()
        if serial_cmd:
            handle_servo_control(serial_cmd)
        time.sleep(0.01)  # 降低CPU占用

# -------------------------- WebSocket逻辑 --------------------------
async def websocket_client():
    global running
    last_heartbeat = 0
    while running:
        try:
            async with connect("ws://new_human_server.qupuba.com:9101/ws") as websocket:
                print("WebSocket连接成功")
                register_msg = json.dumps({"type": "register", "name": board_name})
                await websocket.send(register_msg)
                print(f"发送注册消息: {register_msg}")
                while running:
                    try:
                        message = await asyncio.wait_for(websocket.recv(), timeout=1.0)
                        print(f"收到WebSocket消息: {message}")
                        response = handle_message(message)
                        if response:
                            await websocket.send(response)
                    except asyncio.TimeoutError:
                        if time.time() - last_heartbeat > 15:
                            heartbeat_msg = json.dumps({"type": "heartbeat"})
                            await websocket.send(heartbeat_msg)
                            last_heartbeat = time.time()
                    except exceptions.ConnectionClosed:
                        print("WebSocket连接关闭，重试...")
                        break
        except Exception as e:
            print(f"WebSocket错误: {e}，5秒后重试...")
            await asyncio.sleep(5)

def handle_message(message):
    try:
        data = json.loads(message)
    except json.JSONDecodeError:
        print("WebSocket JSON解析错误")
        return None
    print(f"解析WebSocket数据: {data}")
    msg_type = data.get("type")
    if not msg_type:
        print("缺少type字段")
        return None
    if msg_type == "heartbeat":
        return json.dumps({"type": "heartbeat"})
    elif msg_type == "broadcast":
        print("广播消息:", data.get("content"))
    elif msg_type == "private":
        ext_type = data.get("extType") or data.get("bizType") or data.get("subType")
        if ext_type == "servo_control":
            content = data.get("content")
            try:
                payload = json.loads(content) if isinstance(content, str) else content
                print("private/servo_control:", payload)
                if isinstance(payload, list):
                    for servo in payload:
                        if isinstance(servo, dict):
                            handle_servo_control(servo)
            except Exception as e:
                print("private解析失败:", e)
    elif msg_type == "servo_control":
        content = data.get("content")
        try:
            payload = json.loads(content) if isinstance(content, str) else content
            print("servo_control:", payload)
            if isinstance(payload, list):
                for servo in payload:
                    if isinstance(servo, dict):
                        handle_servo_control(servo)
        except Exception as e:
            print("servo_control解析失败:", e)
    elif msg_type in ("connected", "presence"):
        print("状态消息:", msg_type)
    else:
        print("未知消息类型:", msg_type)
    return None

# -------------------------- 配置加载/保存 --------------------------
def load_config():
    global board_name, wifi_configured
    if os.path.exists(CONFIG_FILE):
        config.read(CONFIG_FILE)
        if "General" in config:
            board_name = config["General"].get("name", "body")
        wifi_configured = True
        return True
    return False

def save_config():
    if not os.path.exists(os.path.dirname(CONFIG_FILE)):
        os.makedirs(os.path.dirname(CONFIG_FILE))
    if "General" not in config:
        config["General"] = {}
    config["General"]["name"] = board_name
    with open(CONFIG_FILE, "w") as f:
        config.write(f)

# -------------------------- 信号处理 --------------------------
def signal_handler(sig, frame):
    global running
    print("\n收到退出信号，关闭程序...")
    running = False
    sys.exit(0)

# -------------------------- 主函数 --------------------------
def main():
    global running
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    if not init_serial():
        print("串口初始化失败，程序退出")
        sys.exit(1)
    if not load_config():
        print("未找到配置文件，使用默认配置")
        save_config()
    print(f"启动舵机控制程序，设备名称: {board_name}")

    file_thread = Thread(target=file_listener_thread, daemon=True)
    file_thread.start()

    serial_thread = Thread(target=serial_listener_thread, daemon=True)
    serial_thread.start()

    try:
        asyncio.get_event_loop().run_until_complete(websocket_client())
    except KeyboardInterrupt:
        print("用户中断程序")
    finally:
        running = False
        if ser and ser.is_open:
            ser.close()
            print("串口已关闭")
        print("程序已退出")

if __name__ == "__main__":
    main()